import os
from types import WrapperDescriptorType
import mindspore
import argparse
from io import BytesIO
import mindspore.dataset as ds
from mindspore.mindrecord import FileWriter
import mindspore.dataset.vision.c_transforms as vision
from PIL import Image

# parse parameters
parser = argparse.ArgumentParser("MindVideo create UCF101 mindrecord script.")
parser.add_argument('--data', type=str, default="", help="Path to the dataset.")
parser.add_argument('--save_path', type=str, default="mindvideo/datasets/UCF101", help="Directory to save the mindrecord file.")
parser.add_argument('-n', '--name', type=str, default="UCF101", help="Mindrecord file name.")
args = parser.parse_args()

MINDRECORD_FILE = args.name
CLS2DIGIT = {}

with open('mindvideo/datasets/UCF101/cls_list.txt', 'r') as f:
    lines = f.readlines()
    for line in lines:
        idx, cls_name = line.split()
        CLS2DIGIT[cls_name]=idx
if os.path.exists(MINDRECORD_FILE):
    os.remove(MINDRECORD_FILE)
    os.remove(MINDRECORD_FILE + ".db")

writer = FileWriter(file_name=MINDRECORD_FILE, shard_num=3)

cv_schema = {"file_name": {"type": "string"}, "label": {"type": "int32"}, "data": {"type": "bytes"}}
writer.add_schema(cv_schema, "it is a UCF101 dataset")

writer.add_index(["file_name", "label"])

video_names = os.listdir(args.data)
for name in video_names:
    video_path = os.path.join(args.data, name)
    video_cls = video_path.split('_')[1]
    with open(video_path, 'rb') as rf:
        content = rf.read()
        writer.write_raw_data([{'file_name': video_path,
                               'label': int(CLS2DIGIT[video_cls]),
                               'data': content}])

writer.commit()
